#include "ps/ps.h"
#include <pybind11/pybind11.h>
#include <pybind11/functional.h>
#include <string>
#include <stdlib.h>
#include <cfloat>

namespace py = pybind11;

template<class T>
void def_kvserver(py::module& m, std::string const& dtype)
{
    std::string const className = "KVServer_" + dtype;
    using Class = ps::KVServer<T>;
    py::class_<Class>(m, className.c_str())
        .def(py::init<int>())
        .def("set_request_handle", &Class::set_request_handle)
        .def("Response", &Class::Response);
}

template<class T>
void def_kvworker(py::module& m, std::string const& dtype)
{
    std::string className = "KVWorker_" + dtype;
    using Class = ps::KVWorker<T>;
    py::class_<Class>(m, className.c_str())
        .def(py::init<int, int>())
        .def("Push", &Class::Push, py::arg("keys"), py::arg("vals"), py::arg("lens") = py::list(), py::arg("cmd")=0, py::arg("cb") = nullptr, py::arg("priority")=0)
        .def("Pull", &Class::Pull, py::arg("keys"), py::arg("vals"), py::arg("lens") = nullptr, py::arg("cmd") = 0, py::arg("cb") = nullptr, py::arg("priority") = 0)
        .def("PushPull", &Class::PushPull, py::arg("keys"), py::arg("vals"), py::arg("outs"), py::arg("lens")=nullptr, py::arg("cmd")=0, py::arg("cb")=nullptr, py::arg("priority")=0)
        .def("Wait", &Class::Wait)
        .def("ZPush", &Class::ZPush, py::arg("keys"), py::arg("vals"), py::arg("lens") = py::list(), py::arg("cmd")=0, py::arg("cb") = nullptr, py::arg("priority")=0)
        .def("ZPull", &Class::ZPull, py::arg("keys"), py::arg("vals"), py::arg("lens") = nullptr, py::arg("cmd") = 0, py::arg("cb") = nullptr, py::arg("priority") = 0)
        .def("ZPushPull", &Class::ZPushPull, py::arg("keys"), py::arg("vals"), py::arg("outs"), py::arg("lens")=nullptr, py::arg("cmd")=0, py::arg("cb")=nullptr, py::arg("priority")=0)
        .def("set_slicer", &Class::set_slicer);

}

template<class T>
void def_SArray(py::module &m, std::string const& dtype)
{
    std::string const className = "SArray_" + dtype;
    using Class = ps::SArray<T>;

    py::class_<Class>(m, className.c_str())
        .def(py::init<::size_t, T>())
        .def(py::init<>())
        .def(py::init<ps::SArray<int32_t>>())
        .def(py::init<ps::SArray<int64_t>>())
        .def(py::init<ps::SArray<float>>())
        .def(py::init<ps::SArray<double>>())
        .def(py::init<std::vector<T> const&>())
        .def("CopyFrom", (void (Class::*)(Class const&))&Class::CopyFrom)
        .def("resize", &Class::resize)
        .def("reserve", &Class::reserve)
        .def("clear", &Class::clear)
        .def("empty", &Class::empty)
        .def("size", &Class::size)
        .def("data", &Class::data, py::return_value_policy::reference)
        .def("back", &Class::back)
        .def("front", &Class::front)
        .def("push_back", &Class::push_back)
        .def("pop_back", &Class::pop_back)
        .def("append", &Class::append)
        .def("segment", &Class::segment)
        .def("__getitem__", [](Class const& c, size_t i){return c[i];})
        .def("__setitem__", [](Class & c, size_t i, T const& v){return c[i] = v;})
        .def("__iter__", [](Class const& c){return py::make_iterator(c.begin(), c.end());});
}

template<class T>
void def_KVPairs(py::module& m, std::string const& dtype)
{
    std::string const className = "KVPairs_" + dtype;
    using Class = ps::KVPairs<T>;

    py::class_<Class>(m, className.c_str())
        .def(py::init<>())
        .def_readwrite("keys", &Class::keys)
        .def_readwrite("vals", &Class::vals)
        .def_readwrite("lens", &Class::lens)
        .def_readwrite("priority", &Class::priority);
}

PYBIND11_MODULE(pyps, m)
{
    m.doc() = "pslite api for ps.h";
    m.def("NumWorkers", &ps::NumWorkers);
    m.def("NumServers", &ps::NumServers);
    m.def("IsWorker", &ps::IsWorker);
    m.def("IsServer", &ps::IsServer);

    m.def("IsScheduler", &ps::IsScheduler);
    m.def("MyRank", &ps::MyRank);

    m.def("Start", &ps::Start);
    m.def("StartAsync", &ps::StartAsync);
    m.def("Finalize", &ps::Finalize);
    m.def("RegisterExitCallback", &ps::RegisterExitCallback);

    def_KVPairs<float>(m, "float32");
    def_KVPairs<double>(m, "float64");
    def_KVPairs<int32_t>(m, "int32");
    def_KVPairs<int64_t>(m, "int64");

    def_kvserver<float>(m, "float32");
    def_kvserver<double>(m, "float64");
    def_kvserver<int32_t>(m, "int32");
    def_kvserver<int64_t>(m, "int64");

    def_SArray<float>(m, "float32");
    def_SArray<double>(m, "float64");
    def_SArray<int32_t>(m, "int32");
    def_SArray<int64_t>(m, "int64");

    def_kvworker<float>(m, "float32");
    def_kvworker<double>(m, "float64");
    def_kvworker<int32_t>(m, "int32");
    def_kvworker<int64_t>(m, "int64_t");

    py::class_<ps::KVMeta>(m, "KVMeta").def(py::init<>())
        .def_readwrite("cmd", &ps::KVMeta::cmd)
        .def_readwrite("push", &ps::KVMeta::push)
        .def_readwrite("pull", &ps::KVMeta::pull)
        .def_readwrite("sender", &ps::KVMeta::sender)
        .def_readwrite("timestamp", &ps::KVMeta::timestamp)
        .def_readwrite("customer_id", &ps::KVMeta::customer_id);

}
